import { StateCreator, StoreMutatorIdentifier } from 'zustand';

type Logger = <
  T,
  Mps extends [StoreMutatorIdentifier, unknown][] = [],
  Mcs extends [StoreMutatorIdentifier, unknown][] = []
>(
  f: StateCreator<T, Mps, Mcs>,
  name?: string
) => StateCreator<T, Mps, Mcs>;

type LoggerImpl = <T>(
  f: StateCreator<T, [], []>,
  name?: string
) => StateCreator<T, [], []>;

const loggerImpl: LoggerImpl = (f, name) => (set, get, store) => {
  const loggedSet: typeof set = (...a) => {
    const before = get();
    set(...a);
    const after = get();
    console.group(name || 'State Update');
    console.log('Previous:', before);
    console.log('Next:', after);
    console.log('Action:', a[0]);
    console.groupEnd();
  };

  return f(loggedSet, get, store);
};

export const logger = loggerImpl as unknown as Logger; 